package task

import (
	aipserver "apedge/apiserver"
	"apedge/dao"
	"apedge/pkg/common"
	"apedge/pkg/configs"
	"apedge/pkg/database"
	"apedge/pkg/loggers"
	"apedge/remote"
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/apulis/go-business/pkg/jobscheduler"
	"github.com/apulis/sdk/go-utils/broker"
	"github.com/apulis/sdk/go-utils/broker/rabbitmq"
	v1 "k8s.io/api/core/v1"
)

var logger = loggers.LogInstance()

func InitMqTask() {
	url := fmt.Sprintf("amqp://%s:%s@%s:%s/",
		configs.Config.Rabbitmq.UserName,
		configs.Config.Rabbitmq.Password,
		configs.Config.Rabbitmq.Host,
		configs.Config.Rabbitmq.Port)
	logger.Infoln("rabbitmq url:", url)

	rabbitMQBroker := rabbitmq.NewBroker(
		broker.Addrs(url),
	)

	if err := rabbitMQBroker.Connect(); err != nil {
		logger.Errorln("rabbitmq1 connect failed!")
		panic(err)
	}

	rabbitMQBroker2 := rabbitmq.NewBroker(
		broker.Addrs(url),
		rabbitmq.ExchangeName(configs.Config.Rabbitmq.ExchangeApsc),
	)

	if err := rabbitMQBroker2.Connect(); err != nil {
		logger.Errorln("rabbitmq2 connect failed!")
		panic(err)
	}

	rabbitMQBroker3 := rabbitmq.NewBroker(
		broker.Addrs(url),
		rabbitmq.ExchangeName(configs.Config.Rabbitmq.ExchangeJs),
	)

	if err := rabbitMQBroker3.Connect(); err != nil {
		logger.Errorln("rabbitmq3 connect failed!")
		panic(err)
	}

	logger.Infoln("rabbitmq connect success!")
	// 监听 jobscheduler
	_, err := rabbitMQBroker3.Subscribe(
		configs.Config.Relevant.ModId,
		RabbitMQMsgHandler,
		rabbitmq.DurableQueue(),
		broker.Queue(configs.Config.Rabbitmq.QueueJs),
	)
	if err != nil {
		panic(err)
	}

	// 监听 模型转换
	_, err = rabbitMQBroker.Subscribe(
		configs.Config.Rabbitmq.TopicModelTransform,
		WorkShopHandler,
		rabbitmq.DurableQueue(),
		broker.Queue(configs.Config.Rabbitmq.QueueModelTransform),
	)
	if err != nil {
		panic(err)
	}

	// 监听 边缘节点上下线
	_, err = rabbitMQBroker2.Subscribe(
		configs.Config.Rabbitmq.TopicApedgeNode,
		ApedgeNodeHandler,
		rabbitmq.DurableQueue(),
		broker.Queue(configs.Config.Rabbitmq.QueueApedgeNode),
	)
	if err != nil {
		panic(err)
	}

	// time.Sleep(time.Duration(2 * time.Second))
	// go TestPub()

}

func RabbitMQMsgHandler(event broker.Event) error {
	var jobMsg jobscheduler.JobMsg
	err := json.Unmarshal(event.Message().Body, &jobMsg)
	if err != nil {
		return err
	}
	logger.Info(fmt.Sprintf("recv infer message: %+v \n", jobMsg))
	err = HandleJobMsg(jobMsg)
	if err != nil {
		return err
	}
	return nil
}

func HandleJobMsg(jobMsg jobscheduler.JobMsg) error {

	// 获取状态
	jobStatus := jobMsg.JobState.GetStatus()
	taskStatus, ok := Job2TaskStatus[jobStatus]
	if !ok {
		return nil
	}

	// 获取多个副本状态
	okPodNum := jobMsg.JobState.DeploymentState.ReadyReplicas
	PodNum := jobMsg.JobState.DeploymentState.Replicas
	repStatus := strconv.Itoa(int(okPodNum)) + "/" + strconv.Itoa(int(PodNum))
	logger.Infoln("pod status:", repStatus)
	tx := database.Db.Begin()
	defer func() {
		if r := recover(); r != nil {
			tx.Rollback()
		}
	}()

	// 更新服务状态
	snt := &dao.ApedgeServiceNode{
		ServiceNodeId: jobMsg.JobId,
	}
	updateMap := make(map[string]interface{})
	updateMap["status"] = taskStatus
	err := snt.UpdateByServiceNodeId(tx, updateMap)
	if err != nil {
		tx.Rollback()
		return err
	}

	return tx.Commit().Error
}

type StudioModelTransformRsp struct {
	ModelId        int    `json:"modelId"`
	ServiceId      string `json:"serviceId"`
	ModelVersionId int    `json:"modelVersionId"`
	Success        bool   `json:"success"`
}

func WorkShopHandler(event broker.Event) error {
	var jobMsg StudioModelTransformRsp

	err := json.Unmarshal(event.Message().Body, &jobMsg)
	if err != nil {
		return err
	}
	logger.Info(fmt.Sprintf("recv workshop rabbitmq message: %+v \n", jobMsg))
	err = HandleWS(jobMsg)
	if err != nil {
		return err
	}
	return nil
}

func HandleWS(jobMsg StudioModelTransformRsp) error {
	var taskStatus string
	// 模型转换判断是否成功
	if jobMsg.Success {
		taskStatus = common.TASK_STATUS_MODEL_TRANSFORM_OK
	} else {
		taskStatus = common.TASK_STATUS_MODEL_TRANSFORM_FAILED
	}

	tx := database.Db.Begin()
	defer func() {
		if r := recover(); r != nil {
			tx.Rollback()
		}
	}()

	task := dao.ApedgeInference{
		ServiceId: jobMsg.ServiceId,
		Status:    taskStatus,
	}
	// err := task.UpdateByServiceId(tx, task)
	// if err != nil {
	// 	tx.Rollback()
	// 	return err
	// }

	item, err := task.GetByServiceId(tx)
	if err != nil {
		tx.Rollback()
		return err
	}
	// 获取服务-节点信息
	sn := dao.ApedgeServiceNode{
		ServiceId: item.ServiceId,
	}
	snList, err := sn.GetListByServiceId(tx)
	if err != nil {
		tx.Rollback()
		return err
	}
	// 成功去启动服务
	if jobMsg.Success {
		for _, snItem := range snList {
			logger.Infoln("start to apflow: ", snItem.ServiceNodeId)
			err = aipserver.StartApflow(*item, snItem)
			if err != nil {
				tx.Rollback()
				return err
			}
			// 更新服务状态
			snt := &dao.ApedgeServiceNode{
				ServiceNodeId: snItem.ServiceNodeId,
			}
			updateMap := make(map[string]interface{})
			updateMap["status"] = common.TASK_STATUS_SCHEDULING
			err = snt.UpdateByServiceNodeId(tx, updateMap)
			if err != nil {
				tx.Rollback()
				return err
			}
		}
	}
	return tx.Commit().Error
}

type NodeInfo struct {
	ID         int64              `json:"id"`
	Action     string             `json:"action"`
	Status     string             `json:"status"`
	Name       string             `json:"name"`
	Roles      string             `json:"roles"`
	Capacity   CapType            `json:"capacity"`
	SystemInfo *v1.NodeSystemInfo `json:"systemInfo"`
}

type CapType struct {
	Cpu    string `json:"cpu"`
	Memory string `json:"memory"`
	Pods   string `json:"pods"`
}

type NodeDeviceMsg struct {
	ID       int64  `json:"id"`
	Type     string `json:"type"`
	Name     string `json:"name"`
	Num      int64  `json:"num"`
	NodeId   int64  `json:"nodeId"`
	NodeName string `json:"nodeName"`
	Action   string `json:"action"`
}

func ApedgeNodeHandler(event broker.Event) error {
	nodeType := event.Message().Header["type"]
	if nodeType == "" {
		logger.Errorln("错误的nodeTypeMsg")
		return errors.New("错误的nodeTypeMsg")
	}
	if nodeType == MSG_TYPE_NODE {
		var jobMsg NodeInfo
		err := json.Unmarshal(event.Message().Body, &jobMsg)
		if err != nil {
			logger.Infoln(err)
			return err
		}
		logger.Info(fmt.Sprintf("recv node message: %+v \n", jobMsg))
		err = HandleNode(jobMsg)
		if err != nil {
			return err
		}
	} else if nodeType == MSG_TYPE_NODE_DEVICE {
		var jobMsg NodeDeviceMsg
		err := json.Unmarshal(event.Message().Body, &jobMsg)
		if err != nil {
			logger.Infoln(err)
			return err
		}
		logger.Info(fmt.Sprintf("recv node message: %+v \n", jobMsg))
		err = HandleNodeDevice(jobMsg)
		if err != nil {
			return err
		}
	} else {
		logger.Errorln("非法的nodeTypeMsg")
	}
	return nil
}

func HandleNode(jobMsg NodeInfo) error {
	// 初始化节点
	if jobMsg.Action == "online" {
		// 打标签 apedgenode={{ nodeId }}
		err := remote.Ks8TagNode(jobMsg.ID, "apedgenode", jobMsg.Name)
		if err != nil {
			logger.Errorln("tag node err: ", err)
		}
		// 打标签 worker=active
		err = remote.Ks8TagNode(jobMsg.ID, "worker", "active")
		if err != nil {
			logger.Errorln("tag node err: ", err)
		}
		// 打标签 archType=?
		arch := "amd64"
		if strings.Contains(jobMsg.SystemInfo.Architecture, "arm") {
			arch = "arm64"
		}
		err = remote.Ks8TagNode(jobMsg.ID, "archType", arch)
		if err != nil {
			logger.Errorln("tag node err: ", err)
		}
		// 打污点 apedgenode=active:NoSchedule
		err = remote.Ks8TaintNode(jobMsg.ID, "apedgenode", "active", "NoExecute")
		if err != nil {
			logger.Errorln("taint node err: ", err)
		}
	}

	tx := database.Db.Begin()
	defer func() {
		if r := recover(); r != nil {
			tx.Rollback()
		}
	}()

	task := dao.ApedgeNode{
		NodeId: jobMsg.Name,
	}
	updateMap := make(map[string]interface{})

	// 节点状态更新操作
	if jobMsg.Action == "updateStatus" {
		if jobMsg.Status == NODE_READY {
			updateMap["status"] = NodeActionToStatus[jobMsg.Status]
			updateMap["arch"] = jobMsg.SystemInfo.Architecture
			updateMap["roles"] = jobMsg.Roles
			updateMap["container_runtime_version"] = jobMsg.SystemInfo.ContainerRuntimeVersion
			updateMap["os_image"] = jobMsg.SystemInfo.OSImage
			cpuNum, err := strconv.ParseInt(jobMsg.Capacity.Cpu, 10, 64)
			if err != nil {
				logger.Errorln("cpuNum bad: ", cpuNum)
				updateMap["cpu"] = 0
				logger.Errorln(err)
			} else {
				updateMap["cpu"] = cpuNum
			}

			memStr := jobMsg.Capacity.Memory
			memStr = strings.Replace(memStr, "Mi", "", -1)
			memNum, err := strconv.ParseInt(memStr, 10, 64)
			if err != nil {
				logger.Errorln("memNum bad: ", memNum)
				updateMap["mem"] = 0
				logger.Errorln(err)
			} else {
				updateMap["mem"] = memNum * 1024 * 1024
			}

		} else if jobMsg.Status == NODE_NOT_READY {
			updateMap["status"] = NodeActionToStatus[jobMsg.Status]
		} else {
			tx.Rollback()
			logger.Errorln("invalid node status!")
			return nil
		}
	} else if jobMsg.Action == "online" {
		updateMap["status"] = NodeActionToStatus[jobMsg.Status]
		updateMap["arch"] = jobMsg.SystemInfo.Architecture
		updateMap["roles"] = jobMsg.Roles
		updateMap["container_runtime_version"] = jobMsg.SystemInfo.ContainerRuntimeVersion
		updateMap["os_image"] = jobMsg.SystemInfo.OSImage
		cpuNum, err := strconv.ParseInt(jobMsg.Capacity.Cpu, 10, 64)
		if err != nil {
			logger.Errorln("cpuNum bad: ", cpuNum)
			updateMap["cpu"] = 0
			logger.Errorln(err)
		} else {
			updateMap["cpu"] = cpuNum
		}

		memStr := jobMsg.Capacity.Memory
		memStr = strings.Replace(memStr, "Mi", "", -1)
		memNum, err := strconv.ParseInt(memStr, 10, 64)
		if err != nil {
			logger.Errorln("memNum bad: ", memNum)
			updateMap["mem"] = 0
			logger.Errorln(err)
		} else {
			updateMap["mem"] = memNum * 1024 * 1024
		}

	} else if jobMsg.Action == "offline" {
		updateMap["status"] = NodeActionToStatus[jobMsg.Status]
	} else {
		tx.Rollback()
		logger.Errorln("invalid node action!")
		return nil
	}
	logger.Infof("updateMap: %+v", updateMap)
	logger.Infof("update nodeId: %s", task.NodeId)
	err := task.UpdateByNodeId(tx, updateMap)
	if err != nil {
		tx.Rollback()
		logger.Errorln("update err: ", err)
		return err
	}

	return tx.Commit().Error
}

func HandleNodeDevice(jobMsg NodeDeviceMsg) error {
	if jobMsg.Action == "online" {
		tx := database.Db.Begin()
		defer func() {
			if r := recover(); r != nil {
				tx.Rollback()
			}
		}()
		nodeTemp := &dao.ApedgeNode{
			NodeId: jobMsg.NodeName,
		}
		updateMap := make(map[string]interface{})
		updateMap["device_spec"] = jobMsg.Name
		err := nodeTemp.UpdateByNodeId(tx, updateMap)
		if err != nil {
			tx.Rollback()
			logger.Errorln("update err: ", err)
			return err
		}
		tx.Commit()
		// 打标签 computeType
		err = remote.Ks8TagNode(jobMsg.NodeId, "computeType", DeviceTagInK8s[jobMsg.Name])
		if err != nil {
			logger.Errorln("tag node err: ", err)
			return err
		}
	}
	return nil
}

func PublishToMq(msg []byte) error {
	url := fmt.Sprintf("amqp://%s:%s@%s:%s/",
		configs.Config.Rabbitmq.UserName,
		configs.Config.Rabbitmq.Password,
		configs.Config.Rabbitmq.Host,
		configs.Config.Rabbitmq.Port)

	rabbitMQBroker := rabbitmq.NewBroker(
		broker.Addrs(url),
		rabbitmq.ExchangeName("default"),
	)
	if err := rabbitMQBroker.Connect(); err != nil {
		logger.Errorln("rabbitmq connect failed! err: ", err)
		return err
	}
	defer rabbitMQBroker.Disconnect()

	err := rabbitMQBroker.Publish(
		"2619",
		&broker.Message{
			Body: msg,
		},
	)
	if err != nil {
		logger.Errorln("publish failed! err: ", err)
		return err
	} else {
		logger.Println("publish success! msg: ", string(msg))
	}
	return nil
}

func TestPub() {
	for {
		// nnnn := "{\"id\":3209,\"action\":\"updateStatus\",\"name\":\"node-1631084077058338876\",\"status\":\"Ready\",\"roles\":\"agent,edge\",\"capacity\":{\"cpu\":\"4\",\"memory\":\"1959Mi\",\"pods\":\"110\"},\"systemInfo\":{\"machineID\":\"\",\"systemUUID\":\"\",\"bootID\":\"\",\"kernelVersion\":\"5.4.0-77-generic\",\"osImage\":\"Ubuntu 18.04.5 LTS\",\"containerRuntimeVersion\":\"docker://20.10.7\",\"kubeletVersion\":\"v1.19.3-kubeedge-v1.5.0\",\"kubeProxyVersion\":\"\",\"operatingSystem\":\"linux\",\"architecture\":\"amd64\"}}"
		err := PublishToMq([]byte("bbbbbbbnnnnnnnnnnnnnnn"))
		if err != nil {
			logger.Errorln("888 ", err)
		}
		logger.Infoln("public ok !")
		time.Sleep(time.Duration(10 * time.Second))

	}
}
